#!/usr/bin/env python3
import csv
import math
import os
from collections import defaultdict

try:
    import numpy as np
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
except Exception as e:
    print('Missing plotting libs:', e)
    np = None
    plt = None

in_path = '/home/ubantu00/zk_demo/build/results.csv'
out_txt = '/home/ubantu00/zk_demo/build/results_summary_full.txt'
out_csv = '/home/ubantu00/zk_demo/build/results_summary_full.csv'
plots_dir = '/home/ubantu00/zk_demo/build/plots'
if not os.path.exists(plots_dir):
    os.makedirs(plots_dir)

groups = defaultdict(list)
with open(in_path, newline='') as f:
    reader = csv.DictReader(f)
    for r in reader:
        key = (r['scenario'], r['type'])
        try:
            r['prove_ms'] = float(r['prove_ms'])
        except:
            r['prove_ms'] = float('nan')
        try:
            r['verify_ms'] = float(r['verify_ms'])
        except:
            r['verify_ms'] = float('nan')
        r['proof_bytes'] = int(r.get('proof_bytes') or 0)
        r['comm_bytes'] = int(r.get('comm_bytes') or 0)
        groups[key].append(r)

summary_rows = []
for (scenario, typ), rows in sorted(groups.items()):
    prove_vals = [r['prove_ms'] for r in rows if not math.isnan(r['prove_ms'])]
    verify_vals = [r['verify_ms'] for r in rows if not math.isnan(r['verify_ms'])]
    proof_bytes_vals = [r['proof_bytes'] for r in rows]
    comm_vals = [r['comm_bytes'] for r in rows]
    n = len(rows)

    def mean(xs):
        return sum(xs)/len(xs) if xs else float('nan')
    def std(xs):
        if len(xs) < 2:
            return 0.0
        m = mean(xs)
        return math.sqrt(sum((x-m)**2 for x in xs)/(len(xs)-1))
    def percentiles(xs, ps):
        if not xs:
            return [float('nan')]*len(ps)
        if np is not None:
            return list(np.percentile(xs, ps))
        # fallback
        xs_sorted = sorted(xs)
        out = []
        for p in ps:
            k = (len(xs_sorted)-1) * (p/100.0)
            f = math.floor(k); c = math.ceil(k)
            if f == c:
                out.append(xs_sorted[int(k)])
            else:
                d0 = xs_sorted[int(f)] * (c-k)
                d1 = xs_sorted[int(c)] * (k-f)
                out.append((d0 + d1))
        return out

    pcts = percentiles(prove_vals, [50,90,99])
    v_pcts = percentiles(verify_vals, [50,90,99])

    row = {
        'scenario': scenario,
        'type': typ,
        'n': n,
        'prove_mean_ms': mean(prove_vals),
        'prove_std_ms': std(prove_vals),
        'prove_median_ms': pcts[0],
        'prove_p90_ms': pcts[1],
        'prove_p99_ms': pcts[2],
        'verify_mean_ms': mean(verify_vals),
        'verify_std_ms': std(verify_vals),
        'verify_median_ms': v_pcts[0],
        'verify_p90_ms': v_pcts[1],
        'verify_p99_ms': v_pcts[2],
        'proof_mean_bytes': mean(proof_bytes_vals),
        'total_comm_bytes': sum(comm_vals)
    }
    summary_rows.append(row)

    # plots
    if plt is not None:
        try:
            plt.clf()
            plt.hist(prove_vals, bins=50)
            plt.title(f'prove_ms distribution {scenario}')
            plt.xlabel('ms')
            plt.ylabel('count')
            plt.tight_layout()
            safe = scenario.replace('/', '_')
            plt.savefig(os.path.join(plots_dir, f'prove_hist_{safe}.png'))

            plt.clf()
            plt.hist(verify_vals, bins=50)
            plt.title(f'verify_ms distribution {scenario}')
            plt.xlabel('ms')
            plt.ylabel('count')
            plt.tight_layout()
            plt.savefig(os.path.join(plots_dir, f'verify_hist_{safe}.png'))

            # boxplot
            plt.clf()
            plt.boxplot([prove_vals, verify_vals], labels=['prove','verify'])
            plt.title(f'prove vs verify boxplot {scenario}')
            plt.tight_layout()
            plt.savefig(os.path.join(plots_dir, f'boxplot_{safe}.png'))
        except Exception as e:
            print('Plot error for', scenario, e)

# write outputs
with open(out_txt, 'w') as f:
    for r in summary_rows:
        f.write(f"Scenario: {r['scenario']} ({r['type']})\n")
        f.write(f"  samples: {r['n']}\n")
        f.write(f"  prove mean(ms): {r['prove_mean_ms']:.6f}, std(ms): {r['prove_std_ms']:.6f}\n")
        f.write(f"  prove median(ms): {r['prove_median_ms']:.6f}, p90: {r['prove_p90_ms']:.6f}, p99: {r['prove_p99_ms']:.6f}\n")
        f.write(f"  verify mean(ms): {r['verify_mean_ms']:.6f}, std(ms): {r['verify_std_ms']:.6f}\n")
        f.write(f"  verify median(ms): {r['verify_median_ms']:.6f}, p90: {r['verify_p90_ms']:.6f}, p99: {r['verify_p99_ms']:.6f}\n")
        f.write(f"  proof mean(bytes): {r['proof_mean_bytes']:.1f}\n")
        f.write(f"  total comm bytes: {r['total_comm_bytes']}\n\n")

import csv as _csv
with open(out_csv, 'w', newline='') as f:
    w = _csv.writer(f)
    w.writerow(['scenario','type','n','prove_mean_ms','prove_std_ms','prove_median_ms','prove_p90_ms','prove_p99_ms','verify_mean_ms','verify_std_ms','verify_median_ms','verify_p90_ms','verify_p99_ms','proof_mean_bytes','total_comm_bytes'])
    for r in summary_rows:
        w.writerow([r['scenario'], r['type'], r['n'],
                    f"{r['prove_mean_ms']:.6f}", f"{r['prove_std_ms']:.6f}", f"{r['prove_median_ms']:.6f}", f"{r['prove_p90_ms']:.6f}", f"{r['prove_p99_ms']:.6f}",
                    f"{r['verify_mean_ms']:.6f}", f"{r['verify_std_ms']:.6f}", f"{r['verify_median_ms']:.6f}", f"{r['verify_p90_ms']:.6f}", f"{r['verify_p99_ms']:.6f}",
                    f"{r['proof_mean_bytes']:.1f}", r['total_comm_bytes']])

print('Wrote', out_txt, 'and', out_csv)
print('Plots (if matplotlib available) in', plots_dir)
